#常规增强方法，镜像、旋转、调光、仿射变换
#为了增加样本的多样性，采用了随即曝光措施
#人脸图像最后统一缩放为128*128像素大小
import cv2
import numpy as np
import os
import shutil
import sys
import random
import time

# 图片增强
def image_augmention(income_path,outcome_path):
    for filenames in os.listdir(income_path):
        img = cv2.imread(original_dir + '/' +filenames)
        if img is None:
            print(filenames + " is error!")
            continue
        # 调整图片的对比度与亮度， 对比度与亮度值都取随机数，这样能增加样本的多样性
        aug_img = relight(img, random.uniform(0.5, 4.5), random.randint(-50, 50))
        cv2.imwrite(outcome_path+'/'+str(time.strftime("%Y%m%d%H%M%S",time.localtime(time.time()))) + '.jpg', aug_img)
        time.sleep(1)
    print("增强完成")
#将增强后的文件移动到原有文件夹
def move_file(old_path, new_path):
    print(old_path)
    print(new_path)
    for filenames in os.listdir(old_path):
        src = os.path.join(old_path,filenames)
        dst = os.path.join(new_path,filenames)
        shutil.move(src,dst)
    print("增强文件已经移动到源文件夹")
# 改变图片的亮度与对比度
def relight(img, light=1, bias=0):
    w = img.shape[1]
    h = img.shape[0]
    #image = []
    for i in range(0,w):
        for j in range(0,h):
            for c in range(3):
                tmp = int(img[j,i,c]*light + bias)
                if tmp > 255:
                    tmp = 255
                elif tmp < 0:
                    tmp = 0
                img[j,i,c] = tmp
    return img

# 对图片进行旋转和翻转
def flip_images(original_filepath,destination_filepath):
    # 对图像进行翻转和特定角度旋转
    for filenames in os.listdir(original_filepath):
        img = cv2.imread(original_filepath + '/' +filenames)

        #水平镜像
        h_flip = cv2.flip(img,1)

        # v_flip = cv2.flip(img,0)#垂直镜像
        # hv_flip = cv2.flip(img,-1)#水平垂直镜像
        
        # 90度旋转
        # rows, cols = img.shape[:2]
        # M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 90, 1)
        # dst_90 = cv2.warpAffine(img, M, (cols, rows))

        # 70度旋转
        # rows, cols = img.shape[:2]
        # M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 70, 1)
        # dst_70 = cv2.warpAffine(img, M, (cols, rows))

        # 60度旋转
        # rows, cols = img.shape[:2]
        # M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 60, 1)
        # dst_60 = cv2.warpAffine(img, M, (cols, rows))

        # 50度旋转
        # rows, cols = img.shape[:2]
        # M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 50, 1)
        # dst_50 = cv2.warpAffine(img, M, (cols, rows))

        # 45度旋转
        # rows, cols = img.shape[:2]
        # M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 45, 1)
        # dst_45 = cv2.warpAffine(img, M, (cols, rows))

        # 40度旋转
        # rows, cols = img.shape[:2]
        # M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 40, 1)
        # dst_40 = cv2.warpAffine(img, M, (cols, rows))

        # 30度旋转
        rows, cols = img.shape[:2]
        M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 30, 1)
        dst_30 = cv2.warpAffine(img, M, (cols, rows))

        # 负15度旋转
        rows, cols = img.shape[:2]
        M = cv2.getRotationMatrix2D((cols / 2, rows / 2), -15, 1)
        dst_negative_30 = cv2.warpAffine(img, M, (cols, rows))

        # 20度旋转
        rows, cols = img.shape[:2]
        M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 20, 1)
        dst_20 = cv2.warpAffine(img, M, (cols, rows))

        # 负12度旋转
        rows, cols = img.shape[:2]
        M = cv2.getRotationMatrix2D((cols / 2, rows / 2), -12, 1)
        dst_negative_20 = cv2.warpAffine(img, M, (cols, rows))

        # 仿射变换，对图像进行变换（三点得到一个变换矩阵）
        # 由三个不同的点确定一个平面，我们可以通过确定三个点的关系来得到转换矩阵
        # 然后再通过warpAffine来进行变换
        point1 = np.float32([[50, 50], [300, 50], [50, 200]])
        point2 = np.float32([[10, 100], [300, 50], [100, 250]])
        M = cv2.getAffineTransform(point1, point2)
        dst1 = cv2.warpAffine(img, M, (cols, rows), borderValue=(102, 112, 122))

        #集中写入
        #镜像写入
        cv2.imwrite(destination_filepath+'/'+str(time.strftime("%Y%m%d%H%M%S",time.localtime(time.time()))) + '.jpg', h_flip)
        time.sleep(1)

        #30度旋转写入
        cv2.imwrite(destination_filepath + "/" + str(time.strftime("%Y%m%d%H%M%S",time.localtime(time.time()))) + '.jpg', dst_30)
        time.sleep(1)

        #-30度旋转写入
        cv2.imwrite(destination_filepath + "/" + str(time.strftime("%Y%m%d%H%M%S",time.localtime(time.time()))) + '.jpg', dst_negative_30)
        time.sleep(1)

        #20度旋转写入
        cv2.imwrite(destination_filepath + "/" + str(time.strftime("%Y%m%d%H%M%S",time.localtime(time.time()))) + '.jpg', dst_20)
        time.sleep(1)

        #-20度旋转写入
        cv2.imwrite(destination_filepath + "/" + str(time.strftime("%Y%m%d%H%M%S",time.localtime(time.time()))) + '.jpg', dst_negative_20)
        time.sleep(1)

        #仿射变换
        cv2.imwrite(destination_filepath + "/" + str(time.strftime("%Y%m%d%H%M%S",time.localtime(time.time()))) + '.jpg', dst1)
        time.sleep(1)

if __name__ == "__main__":
    #读取位置
    original_dir = "source_images"
    # 存储位置
    output_dir = "augmented_images"

    if not os.path.exists(original_dir):#如果不存在文件夹就创建它们
        os.makedirs(original_dir)
    if not os.path.exists(output_dir):#如果不存在文件夹就创建它们
        os.makedirs(output_dir)

    image_augmention(original_dir,output_dir)#将源文件夹中的额文件进行增强
    flip_images(original_dir,output_dir)#旋转
    move_file(output_dir,original_dir)#移动文件

